package com.ai.sitting;

import android.content.Context;
import android.content.res.AssetFileDescriptor;
import android.content.res.AssetManager;

import org.json.JSONException;
import org.json.JSONObject;
import org.tensorflow.lite.Interpreter;

import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.MappedByteBuffer;
import java.nio.channels.FileChannel;

public class PosturePredictor {
    private static final int NUM_SENSORS = 6;
    private final Interpreter tflite;
    private final float[] mean;
    private final float[] scale;

    public PosturePredictor(Context context) throws IOException, JSONException {
        // 加载模型
        tflite = new Interpreter(loadModelFile(context));
        
        // 加载标准化参数
        JSONObject params = loadScalerParams(context);
        mean = new float[NUM_SENSORS];
        scale = new float[NUM_SENSORS];
        
        for (int i = 0; i < NUM_SENSORS; i++) {
            mean[i] = (float) params.getJSONArray("mean").getDouble(i);
            scale[i] = (float) params.getJSONArray("scale").getDouble(i);
        }
    }

    private MappedByteBuffer loadModelFile(Context context) throws IOException {
        String modelPath = "sitting_posture_model.tflite";
        AssetFileDescriptor fileDescriptor = context.getAssets().openFd(modelPath);
        FileInputStream inputStream = new FileInputStream(fileDescriptor.getFileDescriptor());
        FileChannel fileChannel = inputStream.getChannel();
        long startOffset = fileDescriptor.getStartOffset();
        long declaredLength = fileDescriptor.getDeclaredLength();
        return fileChannel.map(FileChannel.MapMode.READ_ONLY, startOffset, declaredLength);
    }

    private JSONObject loadScalerParams(Context context) throws IOException {
        try {
            AssetManager assetManager = context.getAssets();
            InputStream inputStream = assetManager.open("scaler_params.json");
            byte[] buffer = new byte[inputStream.available()];
            inputStream.read(buffer);
            inputStream.close();
            String json = new String(buffer, "UTF-8");
            return new JSONObject(json);
        } catch (Exception e) {
            throw new IOException("Error loading scaler parameters", e);
        }
    }

    public boolean predict(float[] sensorValues) {
        if (sensorValues.length != NUM_SENSORS) {
            throw new IllegalArgumentException("Input must have exactly " + NUM_SENSORS + " values");
        }

        // 标准化数据
        float[] standardizedInput = new float[NUM_SENSORS];
        for (int i = 0; i < NUM_SENSORS; i++) {
            standardizedInput[i] = (sensorValues[i] - mean[i]) / scale[i];
        }

        // 准备输入数据
        float[][] inputArray = new float[1][NUM_SENSORS];
        inputArray[0] = standardizedInput;

        // 准备输出数据
        float[][] outputArray = new float[1][1];

        // 运行推理
        tflite.run(inputArray, outputArray);

        // 返回预测结果
        return outputArray[0][0] >= 0.5f;
    }

    public void close() {
        if (tflite != null) {
            tflite.close();
        }
    }
} 